package org.idea.qiyu.framework.apm.agent.core.jdbc;

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

/**
 * @Author linhao
 * @Date created in 4:37 下午 2021/7/18
 */
public class JdbcAgent {

    public static final String JDBC_CONNECT_CLASS = "com/mysql/jdbc/NonRegisteringDriver";

    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (loader == null || className == null) {
                    return null;
                }
                if (!JDBC_CONNECT_CLASS.equals(className)) {
                    return null;
                }
                return buildJdbcAgentClass(className.replace("/","."));
            }
        },true);
    }

    public static byte[] buildJdbcAgentClass(String className) {
        try {
            ClassPool classPool = ClassPool.getDefault();
            classPool.appendSystemPath();
            CtClass ctClass = classPool.get(className);
            CtMethod ctMethod = ctClass.getDeclaredMethod("connect");
            CtMethod copyMethod = CtNewMethod.copy(ctMethod,ctClass,new ClassMap());
            String newMethodName = ctMethod.getName()+"$Agent";
            ctMethod.setName(newMethodName);
            copyMethod.setBody("{" +
                    "System.out.println(\"connected\");"+
                    "return "+newMethodName+"($$);" +
                    "}");
            ctClass.addMethod(copyMethod);
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}
